# knowledge/base/kb_base.py

import chromadb
import os
import logging
import hashlib
from datetime import datetime
from typing import List, Dict, Optional
from llama_index.core import Settings
from llama_index.embeddings.dashscope import (
    DashScopeEmbedding,
    DashScopeTextEmbeddingModels,
    DashScopeTextEmbeddingType,
)
from llama_index.llms.dashscope import DashScope, DashScopeGenerationModels
from llama_index.core.node_parser import SentenceSplitter
from knowledge.util.tools import logging_helper

class KnowledgeBaseConfig:
    """知识库基础配置类"""
    
    def __init__(self,verbose=False):
        self.logger = logging.getLogger("KnowledgeBase.kb_base")
        self.verbose = verbose
        self._initialized = False
        self.file_registry = {}
    
    def setup_environment(self):
        """设置运行环境 - 优化NLTK设置"""
        try:
            # 预先设置NLTK环境变量，避免重复下载
            self._pre_setup_nltk()
            
            # 尝试导入配置，如果失败则使用简化设置
            try:
                import config
                config.setup_nltk()
                config.setup_ssl()
                config.setup_logging()
                config.setup_base_net()
                logging_helper.info_if_verbose(self.logger,self.verbose,"运行环境设置完成")
            except ImportError:
                self.logger.warning("config模块不存在，使用简化环境设置")
                self._setup_basic_environment()
            except Exception as e:
                self.logger.warning(f"配置环境设置失败，使用简化设置: {e}")
                self._setup_basic_environment()
                
        except Exception as e:
            self.logger.error(f"环境设置失败: {e}")
            raise
    
    def _pre_setup_nltk(self):
        """预先设置NLTK环境变量"""
        import nltk
        
        # 设置NLTK数据目录到用户目录
        nltk_data_dir = os.path.join(os.path.expanduser('~'), 'nltk_data')
        os.makedirs(nltk_data_dir, exist_ok=True)
        
        # 设置环境变量（在导入其他模块之前）
        os.environ['NLTK_DATA'] = nltk_data_dir
        
        # 设置nltk的数据路径（添加到路径开头）
        nltk.data.path.insert(0, nltk_data_dir)
        
        logging_helper.info_if_verbose(self.logger,self.verbose,f"NLTK数据目录预设: {nltk_data_dir}")
    
    def _setup_basic_environment(self):
        """基本环境设置"""
        import logging
        import ssl
        import nltk
        
        # 设置SSL（可选）
        try:
            ssl._create_default_https_context = ssl._create_unverified_context
        except:
            pass
            
        # 设置日志
        logging.basicConfig(
            level=logging.INFO,
            format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
        )
        
        # 检查并下载必要的NLTK数据（只在缺失时）
        try:
            self._ensure_nltk_packages()
        except Exception as e:
            self.logger.warning(f"NLTK包检查失败: {e}")
        
        logging_helper.info_if_verbose(self.logger,self.verbose,"基本环境设置完成")
    
    def _ensure_nltk_packages(self):
        """确保必要的NLTK包存在"""
        import nltk
        
        required_packages = [
            ('tokenizers/punkt', 'punkt'),
            ('tokenizers/punkt_tab', 'punkt_tab'),
        ]
        
        for data_path, package_name in required_packages:
            try:
                nltk.data.find(data_path)
                self.logger.debug(f"NLTK {package_name} 已存在")
            except LookupError:
                self.logger.debug(f"下载 NLTK {package_name}...")
                try:
                    nltk.download(package_name, quiet=True)
                    self.logger.debug(f"NLTK {package_name} 下载完成")
                except Exception as e:
                    self.logger.warning(f"下载 NLTK {package_name} 失败: {e}")
    
    def validate_config(self):
        """验证配置 - 简化版本"""
        try:
            # 只检查必要的API密钥
            api_key = self._get_api_key()
            if not api_key:
                self.logger.error("API密钥未配置")
                return False
            return True
        except Exception as e:
            self.logger.error(f"配置验证异常: {e}")
            return False
    
    def _get_api_key(self):
        """获取API密钥"""
        try:
            import config
            return config.load_api_key()
        except:
            # 从环境变量获取
            return os.getenv('BAILIAN_API_KEY') or os.getenv('DASHSCOPE_API_KEY')

    # 在 initialize_global_settings 方法中增加检索配置
    def initialize_global_settings(self, api_key=None):
        """初始化全局Settings配置"""
        if not api_key:
            api_key = self._get_api_key()
        
        if not api_key:
            raise ValueError("API密钥未配置")
        
        # 1. 设置Embedding模型
        Settings.embed_model = DashScopeEmbedding(
            model_name=DashScopeTextEmbeddingModels.TEXT_EMBEDDING_V2,
            text_type=DashScopeTextEmbeddingType.TEXT_TYPE_DOCUMENT,
            api_key=api_key,
            timeout=30
        )

        # 2. 设置LLM模型
        Settings.llm = DashScope(
            model_name=DashScopeGenerationModels.QWEN_TURBO,
            api_key=api_key,
            temperature=0.1,
            max_tokens=200,
            timeout=10,
            top_p=0.8
        )
        
        # 3. 设置节点解析器
        Settings.node_parser = SentenceSplitter(
            chunk_size=512,
            chunk_overlap=50,
            separator="\n\n"
        )
        
        # 4. 设置全局检索参数
        Settings.chunk_size = 512
        Settings.similarity_top_k = 5
        
        # 5. 设置知识库持久化路径 - 确保使用 enterprise_knowledge
        current_dir = os.path.dirname(os.path.abspath(__file__))
        # 向上两级目录回到 RAG 根目录
        rag_root = os.path.dirname(os.path.dirname(current_dir))
        persist_dir = os.path.join(rag_root, "enterprise_knowledge")  # 正确的路径
        Settings.persist_dir = persist_dir
        
        self._initialized = True
        logging_helper.info_if_verbose(self.logger,self.verbose,f"全局Settings配置完成，持久化路径: {persist_dir}")
    
    def calculate_file_hash(self, file_path: str) -> str:
        """计算文件哈希值"""
        hasher = hashlib.md5()
        try:
            with open(file_path, 'rb') as f:
                buf = f.read()
                hasher.update(buf)
            return hasher.hexdigest()
        except Exception as e:
            self.logger.error(f"计算文件哈希失败 {file_path}: {e}")
            return ""
    
    def load_file_registry(self, registry_path: str = None):
        """加载文件注册表"""
        if not registry_path:
            # 默认放在 knowledge 目录下
            current_dir = os.path.dirname(os.path.abspath(__file__))
            kb_root = os.path.dirname(current_dir)
            registry_path = os.path.join(kb_root, "knowledge_file_registry.json")
            
        import json
        try:
            if os.path.exists(registry_path):
                with open(registry_path, 'r', encoding='utf-8') as f:
                    self.file_registry = json.load(f)
                logging_helper.info_if_verbose(self.logger,self.verbose,f"文件注册表加载完成，包含 {len(self.file_registry)} 个文件记录")
        except Exception as e:
            self.logger.warning(f"加载文件注册表失败: {e}")
    
    def save_file_registry(self, registry_path: str = None):
        """保存文件注册表"""
        if not registry_path:
            # 默认放在 knowledge 目录下
            current_dir = os.path.dirname(os.path.abspath(__file__))
            kb_root = os.path.dirname(current_dir)
            registry_path = os.path.join(kb_root, "knowledge_file_registry.json")
            
        import json
        try:
            with open(registry_path, 'w', encoding='utf-8') as f:
                json.dump(self.file_registry, f, ensure_ascii=False, indent=2)
            logging_helper.info_if_verbose(self.logger,self.verbose,f"文件注册表保存完成，包含 {len(self.file_registry)} 个文件记录")
        except Exception as e:
            self.logger.error(f"保存文件注册表失败: {e}")
    
    def get_chroma_db_path(self):
        """获取ChromaDB路径"""
        return getattr(Settings, "persist_dir", "./knowledge_db")
    
    def is_initialized(self):
        """检查是否已初始化"""
        return self._initialized